import tensorflow as tf


def masked_softmax_cross_entropy(preds, labels, mask):
    """Softmax cross-entropy loss with masking."""
    loss = tf.nn.softmax_cross_entropy_with_logits(logits=preds, labels=labels)
    mask = tf.cast(mask, dtype=tf.float32)
    mask /= tf.reduce_mean(mask)
    loss *= mask

    return tf.reduce_mean(loss)


def masked_accuracy(preds, labels, mask):
    #print("preds:",preds.shape,type(preds),preds)
    #print("labels:",labels.shape,type(labels),labels)
    '''preds: (?, 36) <class 'tensorflow.python.framework.ops.Tensor'> Tensor("graphconvolution_2/SparseTensorDenseMatMul/SparseTensorDenseMatMul:0", shape=(?, 36), dtype=float32)
       labels: (?, 36) <class 'tensorflow.python.framework.ops.Tensor'> Tensor("Placeholder_5:0", shape=(?, 36), dtype=float32)'''
    """Accuracy with masking."""
    correct_prediction = tf.equal(tf.argmax(preds, 1), tf.argmax(labels, 1))
    accuracy_all = tf.cast(correct_prediction, tf.float32)
    mask = tf.cast(mask, dtype=tf.float32)
    mask /= tf.reduce_mean(mask)
    accuracy_all *= mask
    return tf.reduce_mean(accuracy_all)

def dice_coff(preds,labels):
    A = tf.cast(tf.argmax(preds, 1), dtype=tf.float32)  # (10242,)
    B = tf.cast(tf.argmax(labels, 1), dtype=tf.float32)  # (10242,)
    dice = 0.0
    for i in range(1, 36):
        A1 = tf.cast(tf.equal(A, tf.constant(i, dtype=tf.float32)), tf.float32)
        B1 = tf.cast(tf.equal(B, tf.constant(i, dtype=tf.float32)), tf.float32)
        rlt1 = tf.reduce_sum(A1) + tf.reduce_sum(B1)
        rlt2 = tf.reduce_sum(tf.multiply(A1, B1))
        dice += (2.0 * rlt2) / (rlt1 + 0.0001)

    return dice / 31
